from autogen import ConversableAgent
import os

# 配置中文LLM
llm_config = {
    "config_list": [
        {
            "model": "qwen-plus-latest",
            "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "api_key":os.environ["DASHSCOPE_API_KEY"]
        },
        # {"model": "gpt-4", "api_key": os.environ["OPENAI_API_KEY"]}
    ],
}


# The Number Agent always returns the same numbers.
number_agent = ConversableAgent(
    name="Number_Agent",
    system_message="Please return the number sent to you.",
    llm_config=llm_config,
    human_input_mode="NEVER",
)

# The Adder Agent adds 1 to each number it receives.
adder_agent = ConversableAgent(
    name="Adder_Agent",
    system_message="You can only add 1 to the number sent to you and just return the number.",
    llm_config=llm_config,
    human_input_mode="NEVER",
)

# The Subtracter Agent subtracts 1 from each number it receives.
subtracter_agent = ConversableAgent(
    name="Subtracter_Agent",
    system_message="You can only subtract 1 from the number sent to you and just return the number.",
    llm_config=llm_config,
    human_input_mode="NEVER",
)

# The Multiplier Agent multiplies each number it receives by 2.
multiplier_agent = ConversableAgent(
    name="Multiplier_Agent",
    system_message="You can only  multiply the number sent to you by 2 to return the number only.",
    llm_config=llm_config,
    human_input_mode="NEVER",
)

# The Divider Agent divides each number it receives by 2.
divider_agent = ConversableAgent(
    name="Divider_Agent",
    system_message="You can only divide the number sent to you by 2 to return just the number.",
    llm_config=llm_config,
    human_input_mode="NEVER",
)

# The `description` attribute is a string that describes the agent.
# It can also be set in `ConversableAgent` constructor.
adder_agent.description = "Add 1 to the input number."
multiplier_agent.description = "Multiply the input number by 2."
subtracter_agent.description = "Subtract 1 from the input number."
divider_agent.description = "Divide the input number by 2."
number_agent.description = "Returns the input number."

from autogen import GroupChat,GroupChatManager

group_chat = GroupChat(
    agents=[adder_agent, multiplier_agent, subtracter_agent, divider_agent, number_agent],
    messages=[],
    max_round=10,
    speaker_selection_method="auto"
)
group_chat_manager = GroupChatManager(
    groupchat=group_chat,
    llm_config=llm_config
)
chat_result = number_agent.initiate_chat(
    group_chat_manager,
    message="My current number is 7. How do I calculate to get 13 in as few steps as possible? ",
    summary_method="reflection_with_llm",
)
print(chat_result)
